from database.database import text_processing
import numpy as np
import pickle

class rag:
    def __init__(self,question,database_path='./database/data.pkl'):
        self.question=question
        self.path=database_path
        self.set_prompt()

    def binary_search(self,arr, target):
        left, right = 0, len(arr) - 1
        while left <= right:
            mid = (left + right) // 2
            if arr[mid]['id'] == target:
                return mid
            elif arr[mid]['id'] > target:
                right = mid - 1
            else:
                left = mid + 1
        return -1

    def set_score(self):
        with open(self.path, 'rb') as f:
            v_database = pickle.load(f)
        a=text_processing(self.question,bert_path = './bert-base-chinese')
        text_score=[]
        for text in v_database:
            b=text['text_embedding']
            similarity = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
            text_score.append((text['id'],similarity))
        text_score.sort(key = lambda x:x[-1],reverse = True)
        return text_score

    def set_prompt(self):
        self.messages=[
            {"role": "system", "content": "你是一名优秀的音乐作词人，擅长创作中文歌词，请按照用户提出的问题创作中文歌词，歌词不能有违禁脏词，对用户需要有礼貌"},
            {"role": "user", "content": self.question},
            {"role": "assistant", "content": "请告诉我一些相关例子"}
        ]
        score_list=self.set_score()[:3]
        for id,score in score_list:
            result=self.binary_search(v_database,id)
            if result != -1:
                get_model={"role": "user", "content": "例子："}
                get_model['content']+=v_database[result]['text'][:100]+'。'
                self.messages.append(get_model)
